import os
from dotenv import load_dotenv, find_dotenv # 导入 find_dotenv 帮助定位
from langchain.prompts import PromptTemplate
from langchain.chains import (
    LLMChain,
    SimpleSequentialChain,
    TransformChain
)
from langchain.llms import OpenAI

# 加载 .env 文件中的环境变量 (增强调试)
load_dotenv(dotenv_path=find_dotenv(usecwd=True), verbose=True, override=True)

# 从环境变量加载 API 密钥和基础 URL
api_key = os.getenv("OPENAI_API_KEY")
api_base = os.getenv("OPENAI_API_BASE")
os.environ["OPENAI_API_KEY"] = api_key
os.environ["OPENAI_API_BASE"] = api_base

with open("D:/AI Agent/代码和资料/letter.txt") as f:
    letters = f.read()

# 取文件的前3段进行总结
def transform_func(inputs:dict) -> dict:
    text = inputs["text"]
    shortened_text = "\n\n".join(text.split("\n\n")[:3])
    return {"output_text":shortened_text}

#文档转换链
transform_chain = TransformChain(
    input_variables=["text"],
    output_variables=["output_text"],
    transform=transform_func
)

template = """对下面的文字进行总结:
{output_text}

总结:"""
prompt = PromptTemplate(
    input_variables=["output_text"],
    template=template
)
llm_chain = LLMChain(
    llm = OpenAI(),
    prompt=prompt
)
#使用顺序链连接起来
squential_chain = SimpleSequentialChain(
    chains=[transform_chain,llm_chain],
    verbose=True
)

squential_chain.run(letters)